#include "PCGSolver.h"
#include "constants.h"

PCGSolver::PCGSolver()
{
	_clearDataProgram = new GLProgram("Clear Data Proram");
	_clearDataProgram->attach(GL_COMPUTE_SHADER, "shaders/cleardata.glsl");
	_clearDataProgram->link();
	_clearDataProgramUL.n = _clearDataProgram->getUniformLocation("n");
	_clearDataProgramUL.value = _clearDataProgram->getUniformLocation("value");

	GLfloat zeroFloat = 0.f;
	_scalarProductTempBuffer = GLBuffer::newMutable(GL_SHADER_STORAGE_BUFFER, sizeof(GLfloat)*512, NULL, GL_STREAM_DRAW);
	_rBuffer = GLBuffer::newMutable(GL_SHADER_STORAGE_BUFFER, sizeof(GLfloat) * (NUM_NODES + 512), NULL, GL_STREAM_DRAW);
	_zBuffer = GLBuffer::newMutable(GL_SHADER_STORAGE_BUFFER, sizeof(GLfloat) * (NUM_NODES + 512), NULL, GL_STREAM_DRAW);
	_pBuffer = GLBuffer::newMutable(GL_SHADER_STORAGE_BUFFER, sizeof(GLfloat) * (NUM_NODES + 512), NULL, GL_STREAM_DRAW);
	_sBuffer = GLBuffer::newMutable(GL_SHADER_STORAGE_BUFFER, sizeof(GLfloat) * 3, NULL, GL_STREAM_DRAW);

	GLBuffer* buffersToClear[]= {_rBuffer, _zBuffer, _pBuffer};
	for (unsigned i=0;i<3;i++)
	{
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, buffersToClear[i]->name);
		_clearDataProgram->use();
		glUniform1ui(_clearDataProgramUL.n, NUM_NODES + 512);
		glUniform1f(_clearDataProgramUL.value, 0.0f);
		glDispatchCompute((unsigned)ceil((NUM_NODES+512)/256.f), 1, 1);
	}

	GLfloat oneFloat = 1.f;
	_oneBuffer = GLBuffer::newMutable(GL_SHADER_STORAGE_BUFFER, sizeof(GLfloat), &oneFloat, GL_STATIC_DRAW);

	_matrixTimesVectorProgram = new GLProgram("Matrix Times Vector Program");
	_matrixTimesVectorProgram->attach(GL_COMPUTE_SHADER, "shaders/matrixTimesVector.glsl");
	_matrixTimesVectorProgram->link();
	_matrixTimesVectorProgramUL.numRows = _matrixTimesVectorProgram->getUniformLocation("numRows");

	_pcgJacobiPreconditionProgram = new GLProgram("Jacobi Precondition Program");
	_pcgJacobiPreconditionProgram->attach(GL_COMPUTE_SHADER, "shaders/pcgJacobiPrecondition.glsl");
	_pcgJacobiPreconditionProgram->link();
	_pcgJacobiPreconditionProgramUL.numRows = _pcgJacobiPreconditionProgram->getUniformLocation("numRows");
	
	_vectorAddProgram = new GLProgram("Vector Add Program");
	_vectorAddProgram->attach(GL_COMPUTE_SHADER, "shaders/vectorAdd.glsl");
	_vectorAddProgram->link();
	_vectorAddProgramUL.scalarID = _vectorAddProgram->getUniformLocation("scalarID");
	_vectorAddProgramUL.numRows = _vectorAddProgram->getUniformLocation("numRows");
	_vectorAddProgramUL.factor = _vectorAddProgram->getUniformLocation("factor");

	_scalarProduct1Program = new GLProgram("Scalar Product 1 Program");
	_scalarProduct1Program->attach(GL_COMPUTE_SHADER, "shaders/scalarProduct1.glsl");
	_scalarProduct1Program->link();
	_scalarProduct1ProgramUL.n = _scalarProduct1Program->getUniformLocation("n");

	_scalarProduct2Program = new GLProgram("Scalar Product 2 Program");
	_scalarProduct2Program->attach(GL_COMPUTE_SHADER, "shaders/scalarProduct2.glsl");
	_scalarProduct2Program->link();
	_scalarProduct2ProgramUL.n = _scalarProduct2Program->getUniformLocation("n");
	_scalarProduct2ProgramUL.outIndex = _scalarProduct2Program->getUniformLocation("outIndex");
	_scalarProduct2ProgramUL.extraTermIndex = _scalarProduct2Program->getUniformLocation("extraTermIndex");
}

PCGSolver::~PCGSolver()
{
	delete _matrixTimesVectorProgram;
	delete _pcgJacobiPreconditionProgram;
	delete _vectorAddProgram;
	delete _scalarProduct1Program,
	delete _scalarProduct2Program;
	delete _scalarProductTempBuffer;
	delete _rBuffer;
	delete _zBuffer;
	delete _pBuffer;
	delete _sBuffer;
}

void printBuffer(GLBuffer* buffer, unsigned n)
{
	glMemoryBarrier(GL_ALL_BARRIER_BITS);

	glBindBuffer(GL_SHADER_STORAGE_BUFFER, buffer->name);
	
	GLfloat* temp = new GLfloat[n];
	glGetBufferSubData(GL_SHADER_STORAGE_BUFFER, 0, sizeof(GLfloat)*n, temp);
	for (unsigned i=0;i<n;i++)
		printf("%f ",temp[i]);
	printf("\n");
	delete[] temp;
}

void PCGSolver::solve(GLBuffer* initialSolution, unsigned numIterations)
{
	unsigned BLOCK_SIZE = 512;
	unsigned numWorkGroupsSP1 = (unsigned) ceil(_numNodes / (8.f*BLOCK_SIZE));

	//printBuffer(_sparseMatrix, 9*_numNodes);
	//printBuffer(_bVector, _numNodes);

	// comments display state of variables at that time

	glMemoryBarrier(GL_ALL_BARRIER_BITS);

	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, _sparseMatrix->name);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, initialSolution->name);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, _rBuffer->name);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 3, _nodeNeighborhood->name);

	_matrixTimesVectorProgram->use();
	glUniform1ui(_matrixTimesVectorProgramUL.numRows, _numNodes);
	glDispatchCompute((unsigned) ceil(_numNodes / 256.f), 1, 1);

	// r = A*x_0

	glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

	GLuint firstAddBuffers[4] = {
		_bVector->name,
		_rBuffer->name,
		_rBuffer->name,
		_oneBuffer->name
	};

	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, _bVector->name);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, _rBuffer->name);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, _rBuffer->name);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 3, _oneBuffer->name);

	_vectorAddProgram->use();
	glUniform1ui(_vectorAddProgramUL.numRows, _numNodes);
	glUniform1ui(_vectorAddProgramUL.scalarID, 0);
	glUniform1f(_vectorAddProgramUL.factor, -1.f);
	glDispatchCompute((unsigned) ceil(_numNodes / 256.f), 1, 1);

	// r = b-r = b-A*x_0

	//printBuffer(_rBuffer, _numNodes);

	glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

	GLuint precondBuffers[3] = {
		_sparseMatrix->name,
		_rBuffer->name,
		_zBuffer->name
	};
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, _sparseMatrix->name);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, _rBuffer->name);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, _zBuffer->name);

	_pcgJacobiPreconditionProgram->use();
	glUniform1ui(_pcgJacobiPreconditionProgramUL.numRows, _numNodes);
	glDispatchCompute((unsigned) ceil(_numNodes / 256.f), 1, 1);

	glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT | GL_BUFFER_UPDATE_BARRIER_BIT);

	// z = precond(r)

	//printBuffer(_zBuffer, _numNodes);

	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, _zBuffer->name);
	glBindBuffer(GL_COPY_WRITE_BUFFER, _pBuffer->name);
	glCopyBufferSubData(GL_SHADER_STORAGE_BUFFER, GL_COPY_WRITE_BUFFER, 0, 0, sizeof(GLfloat)*_numNodes);

	glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT | GL_BUFFER_UPDATE_BARRIER_BIT);

	// p = z

	//printBuffer(_pBuffer, _numNodes);

	GLuint rzBuffers[4] = {
		_zBuffer->name,
		_rBuffer->name,
		_scalarProductTempBuffer->name,
		_sBuffer->name
	};

	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, _zBuffer->name);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, _rBuffer->name);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, _scalarProductTempBuffer->name);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 3, _sBuffer->name);

	_scalarProduct1Program->use();
	glUniform1ui(_scalarProduct1ProgramUL.n, _numNodes);

	glDispatchCompute(numWorkGroupsSP1,1,1);

	glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

	_scalarProduct2Program->use();
	glUniform1ui(_scalarProduct2ProgramUL.n, numWorkGroupsSP1);
	unsigned sp1OutIndices[3] = {0,2,2};
	glUniform1uiv(_scalarProduct2ProgramUL.outIndex, 3, sp1OutIndices);
	glUniform1ui(_scalarProduct2ProgramUL.extraTermIndex, 0);

	glDispatchCompute(1,1,1);

	// s(1) = <r,z>

	//printBuffer(_sBuffer, 3);

	for (unsigned i=0; i<numIterations; i++)
	{
		
		GLuint secondMultBuffers[4] = {
			_sparseMatrix->name,
			_pBuffer->name,
			_zBuffer->name,
			_nodeNeighborhood->name
		};

		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, _sparseMatrix->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, _pBuffer->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, _zBuffer->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 3, _nodeNeighborhood->name);

		glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);


		_matrixTimesVectorProgram->use();
		glDispatchCompute((unsigned) ceil(_numNodes / 256.f), 1, 1);

		// t = A*p (used z for t)

		//printBuffer(_zBuffer, _numNodes);

		

		GLuint ptBuffers[4] = {
			_pBuffer->name,
			_zBuffer->name,
			_scalarProductTempBuffer->name,
			_sBuffer->name
		};
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, _pBuffer->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, _zBuffer->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, _scalarProductTempBuffer->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 3, _sBuffer->name);

		glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

		_scalarProduct1Program->use();
		glDispatchCompute(numWorkGroupsSP1,1,1);

		glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);
		//begin
		

		_scalarProduct2Program->use();
		unsigned sp2OutIndices[3] = {2,2,1};
		glUniform1uiv(_scalarProduct2ProgramUL.outIndex, 3, sp2OutIndices);

		glDispatchCompute(1,1,1);

		// s(2) = s(1) / <p,t> = <r,z> / <p,A*p>

		//printBuffer(_sBuffer, 3);

		GLuint xAddBuffers[4] = {
			initialSolution->name,
			_pBuffer->name,
			initialSolution->name,
			_sBuffer->name
		};

		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, initialSolution->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, _pBuffer->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, initialSolution->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 3, _sBuffer->name);

		glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

		_vectorAddProgram->use();
		glUniform1ui(_vectorAddProgramUL.scalarID, 1);
		glUniform1f(_vectorAddProgramUL.factor, 1.f);
		glDispatchCompute((unsigned) ceil(_numNodes / 256.f), 1, 1);

		// x = x + s(2)*p

		//printBuffer(initialSolution, _numNodes);

		//if (i == numIterations-1)
		//	break;

		GLuint rAddBuffers[4] = {
			_rBuffer->name,
			_zBuffer->name,
			_rBuffer->name,
			_sBuffer->name
		};

		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, _rBuffer->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, _zBuffer->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, _rBuffer->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 3, _sBuffer->name);

		glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

		glUniform1f(_vectorAddProgramUL.factor, -1.f);
		glDispatchCompute((unsigned) ceil(_numNodes / 256.f), 1, 1);


		//r = r - s(2)*t

		//printBuffer(_rBuffer, _numNodes);

		if (i == numIterations-1)
			break;


		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, _sparseMatrix->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, _rBuffer->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, _zBuffer->name);

		glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

		_pcgJacobiPreconditionProgram->use();
		glDispatchCompute((unsigned) ceil(_numNodes / 256.f), 1, 1);

		// z = precond(r)

		//printBuffer(_zBuffer, _numNodes);

		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, _zBuffer->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, _rBuffer->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, _scalarProductTempBuffer->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 3, _sBuffer->name);

		glMemoryBarrier(GL_TEXTURE_FETCH_BARRIER_BIT | GL_SHADER_STORAGE_BARRIER_BIT);

		_scalarProduct1Program->use();
		glDispatchCompute(numWorkGroupsSP1,1,1);

		glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

		_scalarProduct2Program->use();
		unsigned sp3OutIndices[3] = {0,1,2};
		glUniform1uiv(_scalarProduct2ProgramUL.outIndex, 3, sp3OutIndices);

		glDispatchCompute(1,1,1);

		// [s(1),s(2)} = [<z,r>, <z,r>/s(1)] = [<z,r>, <z,r>/<z_old,r_old>]

		//printBuffer(_sBuffer, 3);

		GLuint pAddBuffers[4] = {
			_zBuffer->name,
			_pBuffer->name,
			_pBuffer->name,
			_sBuffer->name
		};

		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, _zBuffer->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, _pBuffer->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, _pBuffer->name);
		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 3, _sBuffer->name);

		glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

		_vectorAddProgram->use();
		glUniform1f(_vectorAddProgramUL.factor, 1.f);
		glDispatchCompute((unsigned) ceil(_numNodes / 256.f), 1, 1);

		// p = z + s(2)*p

		//printBuffer(_pBuffer, _numNodes);
	}

	glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

	//printf("r:\n");
	//printBuffer(_rBuffer, _numNodes);

	//printf("x:\n");
	//printBuffer(initialSolution, _numNodes);
}

void PCGSolver::set(GLBuffer* sparseMatrix, GLBuffer* nodeNeighborhood, GLBuffer* bVector, unsigned numNodes)
{
	_sparseMatrix = sparseMatrix;
	_nodeNeighborhood = nodeNeighborhood;
	_bVector = bVector;
	_numNodes = numNodes;
}

void PCGSolver::matrixTimesVector(GLBuffer* vector, GLBuffer* solution)
{
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, _sparseMatrix->name);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, vector->name);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, solution->name);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 3, _nodeNeighborhood->name);

	_matrixTimesVectorProgram->use();
	glUniform1ui(_matrixTimesVectorProgramUL.numRows, _numNodes);

	glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

	glDispatchCompute((unsigned) ceil(_numNodes / 256.f), 1, 1);
}

void PCGSolver::scalarProduct(GLBuffer* vector1, GLBuffer* vector2, GLBuffer* solution, unsigned n)
{
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, vector1->name);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, vector2->name);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, _scalarProductTempBuffer->name);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 3, solution->name);

	_scalarProduct1Program->use();
	glUniform1ui(_scalarProduct1ProgramUL.n, n);

	unsigned BLOCK_SIZE = 512;
	unsigned numWorkGroups = (unsigned) ceil(n / (8.f*BLOCK_SIZE));

	glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

	glDispatchCompute(numWorkGroups,1,1);

	/*
	glBindBuffer(GL_SHADER_STORAGE_BUFFER, _scalarProductTempBuffer->name);
	float* temp = new float[512];
	glGetBufferSubData(GL_SHADER_STORAGE_BUFFER, 0, sizeof(GLfloat)*512, temp);
	for (unsigned i=0;i<512;i++)
		printf("%f ", temp[i]);
	printf("\n");
	delete[] temp;
	*/

	glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

	_scalarProduct2Program->use();
	glUniform1ui(_scalarProduct2ProgramUL.n, numWorkGroups);
	glUniform1ui(_scalarProduct2ProgramUL.extraTermIndex, 0);
	GLuint targets[] = {0,2,2};
	glUniform1uiv(_scalarProduct2ProgramUL.outIndex, 3, targets);

	glDispatchCompute(1,1,1);

}